import json
import traceback
from app.jobs.pdfparser.internal.logic.deletetask import PDFDeleteTask
from app.jobs.pdfparser.internal.logic.parsetask import PDFParseTask
from app.jobs.pdfparser.internal.logic.splitter.llm_text_splitter import LlmTextSplitter
from app.jobs.pdfparser.internal.logic.splitter.text_paragraph_splitter import TextParagraphSplitter
from app.jobs.pdfparser.internal.logic.task import PDFTask, PDFTaskType
from app.jobs.pdfparser.internal.server.pdfparserserver import KafkaMsgConsumer
from app.jobs.pdfparser.internal.svc.servicecontext import ServiceContext
from common.utils import logger


class PDFTaskConsumer(KafkaMsgConsumer):
    def __init__(self, svcCtx: ServiceContext):
        super().__init__()
        self.svcCtx = svcCtx
        self._text_splitter = LlmTextSplitter(self.svcCtx.llm, svcCtx)
        
    def consume(self, message) -> None:
        try:
            task = self._parse_msg2task(message)
            task.handle()
        except Exception as e:
            logger.my_logger().error(f'【PDFTaskConsumer】Error consume message: {e}\n{traceback.format_exc()}')
    
    def _parse_msg2task(self, message) -> PDFTask:
        task_info = json.loads(message.value)
        try:
            task_type = PDFTaskType(task_info["task_type"])
        except ValueError as e:
            raise ValueError(f"Invalid task_type: {task_info['task_type']}") from e
        return {
            PDFTaskType.PARSE: PDFParseTask(self.svcCtx, task_info["task_payload"], self._text_splitter),
            PDFTaskType.DELETE: PDFDeleteTask(self.svcCtx, task_info["task_payload"]),
        }[task_type]